"""
This module provides the main function of the iAMT method 
"""

import argparse 
import os
from time import time

import data_handler.data_handler as data_handler
import matching.match_peptides_features as matching
import null_model.null_model as null_model
#import plotting.plot_pfm_data as plot_pfm
from stats.compute_pvalues import DynamicPvalue


##### CONSTANTS 
# used for discretization when calculating the pvalues
STEPS_PER_POSITION = 1000
#####


def load_data(peptide_file, feature_file):
    print "\n--Loading the input data"
    print "Loading %s" % (peptide_file, )
    peptides = data_handler.load_peptide_file(peptide_file)
    print "%d peptides were loaded" % (len(peptides), )

    print "Loading %s" % (feature_file, )
    features = data_handler.load_feature_file(feature_file)        
    print "%d features were loaded" % (len(features), )
    print "--Data loaded."
        
    return peptides, features


def match_peptides(th_peptides, obs_features,  error_tolerance_list, \
         file_root, write_pfms, output_folder):
    print "\n--Matching the peptides to the MS1-features    "
    matching.match_peptides2feat(th_peptides, obs_features, \
            error_tolerance_list)      
    if write_pfms:
        output_file = os.path.join(output_folder, file_root + "-peptide-matched.txt")
        print "PFMs written to %s" % (output_file, )        
        data_handler.write_peptides(th_peptides, output_file)    
    print "--Matching done."    
    

def create_null_model(th_peptides, obs_features, n_bins, \
        error_tolerance_list, file_root, write_null_matches, \
        write_null_probs, output_folder):
    print "\n--Building the null model..."
    null_m, null_probs, null_peptides = \
            null_model.RandRTNullModel.create_null_model(\
            th_peptides, obs_features, n_bins, error_tolerance_list)    
    if write_null_matches:
        output_file = os.path.join(output_folder, \
                file_root + "-null-matches.txt")
        print "Null PFMs written to %s" % (output_file, )                           
        data_handler.write_peptides(null_peptides, output_file)            
    if write_null_probs:
        output_file = os.path.join(output_folder, \
                file_root + "-null-probs.txt")
        print "Null model probabilities written to %s" % (output_file, )                 
        data_handler.write_null_probs(null_probs, output_file)          
    print "--Null model created."
    return null_m                     
           
           
def get_protein_dict(th_peptides, null_m):
    """
    Input:
    th_peptide - list of data_handler.Peptide object with the matched features
                 filled 
    null_m - null model (function that given a Peptide object, return its 
             probability according to the null model)
             
    Output: a list of data_handler.Protein objects              
    """
    
    protein_dict = {}
    for peptide in th_peptides:
        prob = null_m(peptide)
        if len(peptide.matched_features) > 0:
            d = 1
        else:
            d = 0
        for protein in peptide.proteins:
            if protein in protein_dict:
                protein_dict[protein].peptides.append(peptide.sequence)
                protein_dict[protein].probs.append(prob)
                protein_dict[protein].matched.append(d)
            else:
                protein_dict[protein] = data_handler.Protein(protein, \
                        [peptide.sequence], [prob], [d])
    return protein_dict.values()                    
    

def calculate_pvalues(th_peptides, null_m, smooth_pvalue, use_threads):
    print "\n--Computing pvalues, this will take a few minutes..."
    proteins = get_protein_dict(th_peptides, null_m)
    DynamicPvalue.compute_pvalues(proteins, smooth_pvalue, STEPS_PER_POSITION,\
             use_threads)
    print "--P values calculated."   
    return proteins


def write_output_data(proteins, file_root, detailed_out, output_folder):
    print "\n--Saving protein statistics..."
    output_file = os.path.join(output_folder, file_root + "-protein-pvalues.txt")
    proteins.sort(key = lambda protein: len(protein.peptides), reverse = True)
    proteins.sort(key = lambda protein: protein.pvalue)
    data_handler.write_proteins_brief(proteins, output_file)
    print "Final pvalues written to %s" % (output_file, )    
    if detailed_out:
        output_file = os.path.join(output_folder, file_root + \
            "-protein-pvalues-detailed.txt")        
        print "Detailed protein statistics written to %s" % (output_file, )            
        data_handler.write_proteins_detailed(proteins, output_file)
    print "--Done."
        

def run_iAMT(peptide_file, feature_file,  error_tolerance_list, \
        null_model_file, n_bins, smooth_pvalue, use_threads, file_root, \
        write_pfms, write_null_matches, write_null_probs, \
        detailed_out, output_folder):    
    
    # load the input data 
    th_peptides, obs_features = load_data(peptide_file, feature_file)    
    
    # match the peptides to the MS1-features  
    match_peptides(th_peptides, obs_features,  error_tolerance_list, \
             file_root, write_pfms, output_folder)    
    
    # plot mass and rt errors
    # plot_pfm.plot_mass_rt_errors(th_peptides, \
    #    os.path.join(output_folder, file_root + "-errors"))
        
    if null_model_file != None:    
        null_m = null_model.RandRTNullModel.load_null_model(\
                null_model_file)      
    else:
        #create the null model by randomizing retention times 
        null_m = create_null_model(th_peptides, obs_features, n_bins, \
                error_tolerance_list, file_root, write_null_matches,
                write_null_probs, output_folder)
    
    # calculate the pvalues
    proteins = calculate_pvalues(th_peptides, null_m, smooth_pvalue, use_threads)
    
    # write the results to some output file 
    write_output_data(proteins, file_root, detailed_out, output_folder)
    

def string_to_bool(s):
    if s.lower() == 'true' or s.lower() == 't':
        return True
    return False


def print_command_line(args):
    print "############################################"    
    names = ["peptide_file", "feature_file", "null_model_file", \
            "n_bins", "smooth_pvalue", "use_threads", "file_root", \
            "write_pfms", "write_null_probs", "detailed_out", \
            "output_folder"]
    arguments = ["p", "f", "i", "n", "s", "t", "a", "d", "b", "e", "o"]
    if args.r != None:
        names += ["Separate mass and rt errors; mass_tolerance, rt_tolerance"]
        arguments += ["r"]
    else:
        names += ["Combined mass and rt errors; avg_mass_errors, " + \
                "sd_mass_errors, avg_rt_errors, sd_rt_errors, radius"]
        arguments += ["m"]    
    for name, arg in zip(names, arguments):
        if arg == "a" and "a" not in args:
            data = "\"\""
        else:
            data = str(eval("args." + arg))
        print "%s = %s" % (name, data)    
    print "############################################"    
        
        
def main():
    t0 = time()
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', help="File including the in silico peptides. " +\
            "File format: one sequence per line, including the following tab-separated" +\
            " fields: Sequence Predicted_retention_time Theoretical_mass " + \
            "Proteins", metavar="peptide_file", required=True)
    parser.add_argument('-f', help="File including the MS1-features. " + \
            "File format: one feature per line, including the following tab-separated " +\
            "fields: Feature-id Observed_retention_time Observed_mass", \
            metavar = "feature_file", required=True)    
    
    # argument group for using separate retention and mass tolerances 
    sep_group = parser.add_mutually_exclusive_group(required=True)        
    sep_group.add_argument('-r', help="Mass and retention time tolerance" + \
            " windows are treated separately when performing the matching. " +\
            "This argument should be followed by the mass [ppm] and " + \
            "retention time tolerance windows separated by space." + \
            " Only one of the options -r and -m can be used", \
            metavar=("mass_tolerance", "rt_tolerance"), nargs=2, \
            type=float)
    sep_group.add_argument('-m', help="Mass and retention time tolerance" + \
            " windows are combined when performing the matching. " +\
            "This argument should be followed by the average mass error [ppm] "+ \
            "the standard deviation of the mass error, the average retention " + \
            "time error, the standard deviation of the retention time errors and " + \
            "the radius of the accepted circle, space separated. Only one of the options -r " + \
            "and -m can be used ", metavar = ("avg_mass_error", \
            "sd_mass_error", "avg_rt_error", "sd_rt_error", "radius"), \
            nargs=5, type=float)
    parser.add_argument('-i', help="The null model probabilities should " + \
            "be loaded from a file, default: the null model is built by randomizing" + \
            " the predicted retention times of the theoretical peptides", \
            metavar="null_prob_file", required=False, default=None)    
    parser.add_argument('-n', help="Number of bins when randomizing" +\
            " retention times, default = 350", metavar="number_bins", required=False, \
            default=350, type=int)    
    parser.add_argument('-u', help="Write the null model matches " + \
            "to an output file [True/False], default = False",\
            metavar="write_null_matches", \
            required=False, default="False")                    
    parser.add_argument('-s', help="Smooth pvalue [True/False], default = False", \
            metavar="smooth_pvalue", required=False, default="False")        
    parser.add_argument('-t', help="Use threads to calculate pvalues [True/False], default = True", \
            metavar = "threads", required=False, default="True")    
    parser.add_argument('-a', help='File root (all output files will be' + \
            ' prefixed by this). Default: all output files will be prefixed by the word out', \
            metavar="fileroot", required=False, default="out")        
    parser.add_argument('-d', help='Write the peptide-feature matches to an ' + \
            'output file [True/False], default = False', metavar='write_pfms', required=False, 
            default="False")        
    parser.add_argument('-b', help='Write the probabilities according to ' + \
            'the null model to an output file [True/False], default = False', \
            metavar = 'write_null_probs', required=False, default="False")            
    parser.add_argument('-e', help='Write detailed output to an ' + \
            'output file [True/False], default = False', metavar='detailed_out', \
            required = False, default = "False")                
    parser.add_argument('-o', help='Output folder, default = current folder', \
            metavar="output_folder", required=False, default=".")    
    
    args = parser.parse_args()
    
    peptide_file = args.p
    feature_file = args.f
    if args.m != None:        
        error_tolerance_list = args.m
    else:
        error_tolerance_list = args.r
    null_model_file = args.i
    n_bins = args.n
    write_null_matches = string_to_bool(args.u)
    smooth_pvalue = string_to_bool(args.s)
    use_threads = string_to_bool(args.t)
    file_root = args.a
    write_pfms = string_to_bool(args.d)    
    write_null_probs = string_to_bool(args.b)    
    detailed_out = string_to_bool(args.e)
    output_folder = args.o
    
    print_command_line(args)    
    run_iAMT(peptide_file, feature_file, error_tolerance_list, \
            null_model_file, n_bins, smooth_pvalue, use_threads, \
            file_root, write_pfms, write_null_matches, \
            write_null_probs, detailed_out, output_folder)
    
    t1 = time()
    print "\n############################################"  
    print "Analysis time: %.1f minutes" % ((t1 - t0) / 60)
    print "############################################"    
    
if __name__ == '__main__':
    main()
